from langchain.document_loaders import TextLoader, UnstructuredFileLoader
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import Chroma
import sentence_transformers

# 导入文本
loader = TextLoader("../压铸机2/压铸机2.txt", "utf-8")
# 将文本转成 Document 对象
data = loader.load()
print(f'documents:{len(data)}')

# 初始化加载器
text_splitter = RecursiveCharacterTextSplitter(chunk_size=100, chunk_overlap=0)
# 切割加载的 document
split_docs = text_splitter.split_documents(data)
print("split_docs size:", len(split_docs))

embedding_model_dict = {
    "text2vec": "C:\\Work\\llm\\text2vec-large-chinese"
}

EMBEDDING_MODEL = "text2vec"
# 初始化 hugginFace 的 embeddings 对象
embeddings = HuggingFaceEmbeddings(model_name=embedding_model_dict[EMBEDDING_MODEL])
embeddings.client = sentence_transformers.SentenceTransformer(
    embeddings.model_name, device='cuda')



# 初始化加载器
db = Chroma.from_documents(split_docs, embeddings, persist_directory="./chroma/casting2")
# # 持久化
db.persist()